{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a192dc3a",
   "metadata": {
    "id": "TKvYiJgnYExi"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/clutter.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1410a65",
   "metadata": {
    "id": "A4QOaw_zYLfI",
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from IPython.display import clear_output\n",
    "from pydrake.all import (\n",
    "    AbstractValue,\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Box,\n",
    "    Capsule,\n",
    "    ContactResults,\n",
    "    ContactVisualizer,\n",
    "    ContactVisualizerParams,\n",
    "    CoulombFriction,\n",
    "    Cylinder,\n",
    "    DiagramBuilder,\n",
    "    Ellipsoid,\n",
    "    FixedOffsetFrame,\n",
    "    JointSliders,\n",
    "    LeafSystem,\n",
    "    MeshcatVisualizer,\n",
    "    Parser,\n",
    "    PlanarJoint,\n",
    "    Rgba,\n",
    "    RigidTransform,\n",
    "    RotationMatrix,\n",
    "    SpatialInertia,\n",
    "    Sphere,\n",
    "    StartMeshcat,\n",
    "    UnitInertia,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.scenarios import AddShape\n",
    "from manipulation.utils import ConfigureParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22e9071d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50154803",
   "metadata": {
    "id": "7aGJ9e3lJlF8"
   },
   "source": [
    "# Contact force \"inspector\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c32e2cb",
   "metadata": {
    "id": "7dimkJhd9DjP"
   },
   "outputs": [],
   "source": [
    "class PrintContactResults(LeafSystem):\n",
    "    def __init__(self):\n",
    "        LeafSystem.__init__(self)\n",
    "        self.DeclareAbstractInputPort(\n",
    "            \"contact_results\", AbstractValue.Make(ContactResults())\n",
    "        )\n",
    "        self.DeclareForcedPublishEvent(self.Publish)\n",
    "        self.DeclarePeriodicPublishEvent(0.1, 0, self.Publish)\n",
    "\n",
    "    def Publish(self, context):\n",
    "        formatter = {\"float\": lambda x: \"{:5.2f}\".format(x)}\n",
    "        results = self.get_input_port().Eval(context)\n",
    "\n",
    "        if results.num_hydroelastic_contacts() == 0:\n",
    "            print(\"no contact\")\n",
    "        for i in range(results.num_hydroelastic_contacts()):\n",
    "            info = results.hydroelastic_contact_info(i)\n",
    "            print(f\"F_Ac_W = {info.F_Ac_W()}\")\n",
    "\n",
    "        clear_output(wait=True)\n",
    "\n",
    "\n",
    "def contact_force_inspector(slope=0.0, mu=1.0, second_brick=False):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.01)\n",
    "\n",
    "    AddShape(\n",
    "        plant,\n",
    "        Box(10.0, 10.0, 10.0),\n",
    "        \"ground\",\n",
    "        mass=1,\n",
    "        mu=1,\n",
    "        color=[0.9, 0.9, 0.9, 0.8],\n",
    "    )\n",
    "    X_WBox = RigidTransform(RotationMatrix.MakeYRotation(slope), [0, 0, -5.05])\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(), plant.GetFrameByName(\"ground\"), X_WBox\n",
    "    )\n",
    "\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    parser.AddModelsFromUrl(\"package://manipulation/hydro/061_foam_brick.sdf\")\n",
    "    frame = plant.AddFrame(\n",
    "        FixedOffsetFrame(\n",
    "            \"planar_joint_frame\",\n",
    "            plant.world_frame(),\n",
    "            RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2)),\n",
    "        )\n",
    "    )\n",
    "    plant.AddJoint(\n",
    "        PlanarJoint(\n",
    "            \"brick\",\n",
    "            frame,\n",
    "            plant.GetFrameByName(\"base_link\"),\n",
    "            damping=[0, 0, 0],\n",
    "        )\n",
    "    )\n",
    "\n",
    "    if second_brick:\n",
    "        parser.SetAutoRenaming(True)\n",
    "        brick2 = parser.AddModelsFromUrl(\n",
    "            \"package://manipulation/hydro/061_foam_brick.sdf\"\n",
    "        )[0]\n",
    "        plant.AddJoint(\n",
    "            PlanarJoint(\n",
    "                \"brick2\",\n",
    "                frame,\n",
    "                plant.GetFrameByName(\"base_link\", brick2),\n",
    "                damping=[0, 0, 0],\n",
    "            )\n",
    "        )\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    meshcat.Set2dRenderMode(xmin=-0.2, xmax=0.2, ymin=-0.2, ymax=0.3)\n",
    "\n",
    "    cparams = ContactVisualizerParams()\n",
    "    cparams.force_threshold = 1e-6\n",
    "    cparams.newtons_per_meter = 1.0\n",
    "    cparams.radius = 0.002\n",
    "    contact_visualizer = ContactVisualizer.AddToBuilder(\n",
    "        builder, plant, meshcat, cparams\n",
    "    )\n",
    "\n",
    "    print_contact_results = builder.AddSystem(PrintContactResults())\n",
    "    builder.Connect(\n",
    "        plant.get_contact_results_output_port(),\n",
    "        print_contact_results.get_input_port(),\n",
    "    )\n",
    "\n",
    "    lower_limit = [-0.1, -0.1, -np.pi / 2.0]\n",
    "    upper_limit = [0.1, 0.1, np.pi / 2.0]\n",
    "    q0 = [0, 0, 0]\n",
    "    if second_brick:\n",
    "        lower_limit += lower_limit\n",
    "        upper_limit += upper_limit\n",
    "        q0 += [0.07, 0.07, 0.0]\n",
    "\n",
    "    default_interactive_timeout = None if running_as_notebook else 1.0\n",
    "    sliders = builder.AddSystem(\n",
    "        JointSliders(\n",
    "            meshcat,\n",
    "            plant,\n",
    "            initial_value=q0,\n",
    "            lower_limit=lower_limit,\n",
    "            upper_limit=upper_limit,\n",
    "            step=0.001,\n",
    "        )\n",
    "    )\n",
    "    diagram = builder.Build()\n",
    "    sliders.Run(diagram, default_interactive_timeout)\n",
    "    meshcat.DeleteAddedControls()\n",
    "\n",
    "\n",
    "contact_force_inspector(slope=0.1, mu=0.5, second_brick=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "886ff489",
   "metadata": {
    "id": "2ctYY0SHJlGB"
   },
   "source": [
    "# Contact results \"inspector\"\n",
    "\n",
    "This simple visualization shows some of the complexity of the contact geometry problem.  I will make it better, but right now, when you move the objects into contact of each other you will see three points:  the contact point is in **red**, the contact normal is added to the contact point with the tip as **green**, and the (scaled) contact force tip is drawn in **blue**.  Contact points on the bodies are drawn in **orange**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5827dc3f",
   "metadata": {
    "id": "648WXpNpJlGF"
   },
   "outputs": [],
   "source": [
    "shapes = {\n",
    "    \"Point\": Sphere(0.01),\n",
    "    \"Sphere\": Sphere(1.0),\n",
    "    \"Cylinder\": Cylinder(1.0, 2.0),\n",
    "    \"Box\": Box(1.0, 2.0, 3.0),\n",
    "    \"Capsule\": Capsule(1.0, 2.0),\n",
    "    \"Ellipsoid\": Ellipsoid(1.0, 2.0, 3.0),\n",
    "}\n",
    "\n",
    "\n",
    "def contact_inspector(shape_name_A, shape_name_B):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "    AddShape(\n",
    "        plant,\n",
    "        shapes[shape_name_A],\n",
    "        \"A\",\n",
    "        mass=1,\n",
    "        mu=1,\n",
    "        color=[0.9, 0.5, 0.5, 0.5],\n",
    "    )\n",
    "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"A\"))\n",
    "\n",
    "    AddShape(\n",
    "        plant,\n",
    "        shapes[shape_name_B],\n",
    "        \"B\",\n",
    "        mass=1,\n",
    "        mu=1,\n",
    "        color=[0.5, 0.5, 0.9, 0.5],\n",
    "    )\n",
    "    frame = plant.AddFrame(\n",
    "        FixedOffsetFrame(\n",
    "            \"planar_joint_frame\",\n",
    "            plant.world_frame(),\n",
    "            RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2)),\n",
    "        )\n",
    "    )\n",
    "    plant.AddJoint(\n",
    "        PlanarJoint(\"B\", frame, plant.GetFrameByName(\"B\"), damping=[0, 0, 0])\n",
    "    )\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    meshcat.Set2dRenderMode(xmin=-3.0, xmax=3.0, ymin=-0.2, ymax=3.0)\n",
    "\n",
    "    cparams = ContactVisualizerParams()\n",
    "    cparams.force_threshold = 1e-6\n",
    "    cparams.newtons_per_meter = 1.0\n",
    "    cparams.radius = 0.002\n",
    "    contact_visualizer = ContactVisualizer.AddToBuilder(\n",
    "        builder, plant, meshcat, cparams\n",
    "    )\n",
    "\n",
    "    print_contact_results = builder.AddSystem(PrintContactResults())\n",
    "    builder.Connect(\n",
    "        plant.get_contact_results_output_port(),\n",
    "        print_contact_results.get_input_port(),\n",
    "    )\n",
    "\n",
    "    lower_limit = [-3, -3, -np.pi / 2.0]\n",
    "    upper_limit = [3, 3, np.pi / 2.0]\n",
    "    q0 = [1.2, 1.2, 0.0]\n",
    "\n",
    "    default_interactive_timeout = None if running_as_notebook else 1.0\n",
    "    sliders = builder.AddSystem(\n",
    "        JointSliders(\n",
    "            meshcat,\n",
    "            plant,\n",
    "            initial_value=q0,\n",
    "            lower_limit=lower_limit,\n",
    "            upper_limit=upper_limit,\n",
    "        )\n",
    "    )\n",
    "    diagram = builder.Build()\n",
    "    sliders.Run(diagram, default_interactive_timeout)\n",
    "    meshcat.DeleteAddedControls()\n",
    "\n",
    "\n",
    "contact_inspector(\"Box\", \"Sphere\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42269089",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
